#include "angle_estimator.h"

// 波束形成 - 计算空间谱
std::vector<float> AngleEstimator::beamforming(const std::vector<Complex> &antenna_data,
                                               const std::vector<float> &angles)
{
    std::vector<float> spatial_spectrum(angles.size(), 0.0f);

    for (size_t i = 0; i < angles.size(); ++i)
    {
        float angle_rad = angles[i] * M_PI / 180.0f;
        Complex sum(0, 0);

        for (int j = 0; j < params.num_antennas; ++j)
        {
            // 计算波程差引起的相位差
            float phase_shift = 2.0f * M_PI * antenna_positions[j] *
                                std::sin(angle_rad) / params.light_speed *
                                params.carrier_frequency;

            // 应用相位补偿
            sum += antenna_data[j] * std::exp(Complex(0, -phase_shift));
        }

        spatial_spectrum[i] = std::norm(sum);
    }

    return spatial_spectrum;
}

// MUSIC算法 - 高分辨率测角
std::vector<float> AngleEstimator::musicAlgorithm(const std::vector<Complex> &antenna_data,
                                                  const std::vector<float> &angles, int signal_num)
{
    int snapshot_num = 1; // 假设单次快拍
    MatrixXcf data_matrix(params.num_antennas, snapshot_num);

    // 填充数据矩阵
    for (int i = 0; i < params.num_antennas; ++i)
    {
        data_matrix(i, 0) = antenna_data[i];
    }

    // 计算协方差矩阵
    MatrixXcf covariance = data_matrix * data_matrix.adjoint() / snapshot_num;

    // 特征值分解
    Eigen::SelfAdjointEigenSolver<MatrixXcf> eigensolver(covariance);
    if (eigensolver.info() != Eigen::Success)
    {
        return std::vector<float>(angles.size(), 0.0f);
    }

    // 获取特征值和特征向量
    VectorXcf eigenvalues = eigensolver.eigenvalues();
    MatrixXcf eigenvectors = eigensolver.eigenvectors();

    // 排序特征值 (降序)
    std::vector<int> indices(params.num_antennas);
    for (int i = 0; i < params.num_antennas; ++i)
        indices[i] = i;

    std::sort(indices.begin(), indices.end(),
              [&eigenvalues](int a, int b)
              {
                  return std::real(eigenvalues[a]) > std::real(eigenvalues[b]);
              });

    // 构建噪声子空间
    MatrixXcf noise_subspace(params.num_antennas, params.num_antennas - signal_num);
    for (int i = 0; i < params.num_antennas - signal_num; ++i)
    {
        noise_subspace.col(i) = eigenvectors.col(indices[signal_num + i]);
    }

    // 计算MUSIC谱
    std::vector<float> music_spectrum(angles.size(), 0.0f);

    for (size_t i = 0; i < angles.size(); ++i)
    {
        float angle_rad = angles[i] * M_PI / 180.0f;
        VectorXcf steering_vector(params.num_antennas);

        // 构建导向矢量
        for (int j = 0; j < params.num_antennas; ++j)
        {
            float phase_shift = 2.0f * M_PI * antenna_positions[j] *
                                std::sin(angle_rad) / params.light_speed *
                                params.carrier_frequency;
            steering_vector(j) = std::exp(Complex(0, -phase_shift));
        }

        // 计算MUSIC谱
        Complex denominator = steering_vector.adjoint() * noise_subspace *
                              noise_subspace.adjoint() * steering_vector;
        music_spectrum[i] = 1.0f / std::abs(denominator);
    }

    return music_spectrum;
}

// 估计角度
float AngleEstimator::estimateAngle(const std::vector<Complex> &antenna_data,
                                    const std::vector<float> &angles)
{
    auto spectrum = beamforming(antenna_data, angles);

    // 找到峰值位置
    auto max_it = std::max_element(spectrum.begin(), spectrum.end());
    int peak_idx = std::distance(spectrum.begin(), max_it);

    return angles[peak_idx];
}